import os
import json
import datetime
import traceback
import shutil
from preprocess import trans2txt, doc2text
from ltp import LTP


headers = ['Criminals', 'Gender', 'Ethnicity', 'Birthplace', 'Accusation', 'Courts']
stop_words = []
accusations = []
all_infos = []
cur_info = []
labels = {
    'noun': [],
    'verb': [],
    'adj': []
}

ltp = LTP()
ltp.init_dict('dataset/user_dict.txt')


def getStopWords():
    if not len(stop_words) == 0:
        return
    with open('dataset/stop_word.txt', 'r', encoding='utf-8') as f:
        for word in f.readlines():
            stop_words.append(word.strip())


def getAccusations():
    if not len(accusations) == 0:
        return
    with open('dataset/accusations.txt', 'r', encoding='utf-8') as f:
        for acc in f.readlines():
            accusations.append(acc.strip())


def massExtract(path):
    if os.path.isfile(path):
        trans2text(path)
        new_filename = 'dataset/text/' + path[0:path.rfind('.')] + '.txt'
        print(new_filename)
        extract(path)
        return
    for f in os.listdir(path):
        fname = path + '/' + f
        if not (os.path.isfile(fname) and trans2text(fname)):
            continue
        try:
            new_filename = 'dataset/text/' + f[0:f.rfind('.')] + '.txt'
            print(new_filename)
            extract(new_filename)
            pass
        except:
            traceback.print_exc()
            print('\033[1;35m Error \033[0m!')


class FoundFlag:
    def __init__(self):
        self.name = False
        self.gender = False
        self.ethnicity = False
        self.birthplace = False
        self.accusations = False
        self.courts = False


acc_not_found = False


def extract(filename):
    genders = ['男', '女']
    criminals, gender, ethnicity, birthplace, accusation, courts = [], [], [], [], [], []
    ac, ag, ae, ab, aa, acs = [], [], [], [], [], []
    global acc_not_found
    acc_not_found = False

    def findCriminal(cur_sen, cur_pos, i):
        # if cur_pos[i + 1] in ['nh', 'n']:
        if cur_pos[i + 1] in ['nh'] or cur_sen[i + 1].endswith('x') or cur_sen[i + 1].endswith('某') or cur_sen[
            i + 1].endswith('×'):
            if 4 < len(cur_sen[i + 1]) or len(cur_sen[i + 1]) < 2:
                return False
            for c in criminals:
                if c in cur_sen[i + 1]:
                    return False
            criminals.append(cur_sen[i + 1])
            if cur_sen[i + 1] not in labels['noun']:
                labels['noun'].append(cur_sen[i + 1])
            return True
        return False

    def isProperNoun(p, s):
        if len(s) < 2:
            return False
        if p in ['n', 'nh', 'ni', 'nl', 'ns', 'nz']:
            if s.endswith('族'):
                return True

    def findNoun(cur_sen, cur_pos):
        idx = 0
        while idx < len(cur_sen):
            if cur_pos[idx] == 'ns':
                i = idx
                place = ''
                while i < len(cur_sen) and (
                        cur_pos[i] == 'ns' or cur_sen[i][-1] in ['市', '区', '县', '省', '村', '镇', '乡']):
                    place += cur_sen[i]
                    i += 1
                if place != '':
                    if place not in labels['noun']:
                        labels['noun'].append(place)
                    idx = i
            elif (cur_sen[idx].endswith('纠纷') or cur_sen[idx].endswith('罪')) \
                    and cur_sen[idx] in accusations:
                acc = []
                i = idx
                while i < len(cur_sen) and cur_sen[i] in accusations and (
                        cur_sen[i].endswith('纠纷') or cur_sen[i].endswith('罪')):
                    acc.append(cur_sen[i])
                    i += 1
                idx = i - 1
                if acc not in labels['noun']:
                    labels['noun'].append('，'.join(acc))
            idx += 1

    def getWords(ls):
        global acc_not_found
        if ls == 1 and seg[0][-1].endswith('法院'):
            court = ''.join(seg[0])
            print('court', court)
            if court not in courts:
                courts.append(court)
                labels['noun'].append(court)
            return True
        flag = FoundFlag()
        for i in range(len(seg)):
            cur_sen, cur_pos = [], []
            for j in range(len(seg[i])):
                if pos[i][j] == 'wp' or seg[i][j] in stop_words:
                    continue
                cur_sen.append(seg[i][j])
                cur_pos.append(pos[i][j])
                if isProperNoun(pos[i][j], seg[i][j]) or seg[i][j] in genders:
                    if len(seg[i][j]) > 1 and seg[i][j] not in labels['noun']:
                        labels['noun'].append(seg[i][j])
                elif pos[i][j] in ['v']:
                    if len(seg[i][j]) > 1 and seg[i][j] not in labels['verb']:
                        labels['verb'].append(seg[i][j])
                elif pos[i][j] in ['a']:
                    if len(seg[i][j]) > 1 and seg[i][j] not in labels['adj']:
                        labels['adj'].append(seg[i][j])
                # print('{}/{}'.format(seg[i][j], pos[i][j]))
            # print()

            findNoun(cur_sen, cur_pos)

            idx = 0
            while idx < len(cur_sen):
                if cur_sen[idx].endswith('法院') and cur_pos[idx] == 'ns':
                    if cur_sen[idx] not in courts:
                        courts.append(cur_sen[idx])
                        labels['noun'].append(cur_sen[idx])
                        print('court:', cur_sen[idx])
                elif cur_sen[idx].endswith('分院') and cur_pos[idx] == 'ns':
                    if cur_sen[idx] not in courts:
                        courts.append(cur_sen[idx])
                        labels['noun'].append(cur_sen[idx])
                        print('court:', cur_sen[idx])
                elif not flag.name and cur_sen[idx] == '被告人':
                    if idx >= len(cur_sen) - 1:
                        break
                    flag.name = findCriminal(cur_sen, cur_pos, idx)
                    if flag.name:
                        idx += 1
                        print('criminal:', cur_sen[idx])
                elif flag.name:
                    if not flag.gender and cur_sen[idx] in genders:
                        gender.append(cur_sen[idx])
                        flag.gender = True
                        print('gender:', cur_sen[idx])
                    elif not flag.ethnicity and cur_sen[idx].endswith('族'):
                        ethnicity.append(cur_sen[idx])
                        flag.ethnicity = True
                        print('ethnicity:', cur_sen[idx])
                    elif not flag.birthplace and cur_sen[idx] in ['户籍地', '出生于', '住']:
                        i = idx
                        i += 1
                        place = ''
                        while i < len(cur_sen) and cur_pos[i] == 'ns':
                            place += cur_sen[i]
                            i += 1
                        if place != '':
                            birthplace.append(place)
                            flag.birthplace = True
                            idx = i - 1
                            print('birthplace:', place)
                    elif not flag.birthplace and cur_pos[idx] == 'ns':
                        i = idx
                        place = ''
                        while i < len(cur_sen) and (
                                cur_pos[i] == 'ns' or cur_sen[i][-1] in ['市', '区', '县', '省', '村', '镇', '乡']):
                            place += cur_sen[i]
                            i += 1
                        if place != '' and i < len(cur_sen) and cur_sen[i] == '人':
                            birthplace.append(place)
                            flag.birthplace = True
                            idx = i
                            print('birthplace:', place)
                    elif not flag.accusations \
                            and (cur_sen[idx].endswith('纠纷') or cur_sen[idx].endswith('罪')) \
                            and cur_sen[idx] in accusations:
                        acc = []
                        i = idx
                        while i < len(cur_sen) and cur_sen[i] in accusations and (
                                cur_sen[i].endswith('纠纷') or cur_sen[i].endswith('罪')):
                            acc.append(cur_sen[i])
                            i += 1
                        idx = i - 1
                        accusation.append('，'.join(acc))
                        print('accusation:', acc)
                        flag.accusations = True
                idx += 1
        if flag.name and (not flag.ethnicity or not flag.birthplace or not flag.gender or not flag.accusations):
            print("error")
            print(flag.name, flag.gender, flag.ethnicity, flag.birthplace, flag.accusations)
            if not flag.accusations:
                c = criminals.pop()
                ac.append(c)
                if flag.gender:
                    g = gender.pop()
                    ag.append(g)
                else:
                    ag.append('')
                if flag.ethnicity:
                    e = ethnicity.pop()
                    ae.append(e)
                else:
                    ae.append('')
                if flag.birthplace:
                    b = birthplace.pop()
                    ab.append(b)
                else:
                    ab.append('')
                acc_not_found = True
                return False
            elif acc_not_found:
                if not flag.gender:
                    gender.append('未知')
                if not flag.ethnicity:
                    ethnicity.append('未知')
                if not flag.birthplace:
                    birthplace.append('未知')
                index = 0
                for c in ac:
                    if c in criminals:
                        idx = criminals.index(c)
                        if ag[index] != '':
                            gender[idx] = ag[index]
                            flag.gender = True
                        if ae[index] != '':
                            ethnicity[idx] = ae[index]
                            flag.ethnicity = True
                        if ab[index] != '':
                            birthplace[idx] = ab[index]
                            flag.birthplace = True
                    index += 1
                if len(ac) == 0:
                    acc_not_found = False
            else:
                if not flag.gender:
                    gender.append('未知')
                if not flag.ethnicity:
                    ethnicity.append('未知')
                if not flag.birthplace:
                    birthplace.append('未知')
            return False
        return True

    def printAll():
        if len(criminals) == 0:
            with open('dataset/error.txt', 'a', encoding='utf-8') as ef:
                ef.write(filename.split('/')[-1].split('.')[0] + '\n')
        global cur_info
        cur_info = []
        for i in range(len(criminals)):
            infos = [criminals[i], gender[i], ethnicity[i], birthplace[i], accusation[i], '，'.join(courts)]
            info_dict = dict(zip(headers, infos))
            info_dict['File'] = filename.split('/')[-1].split('.')[0]
            all_infos.append(info_dict)
            cur_info.append(info_dict)
            print(info_dict)

    with open(filename, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            if line.strip() == '':
                continue
            sents = ltp.sent_split([line.strip()])
            try:
                seg, hidden = ltp.seg(sents)
            except RuntimeError:
                continue
            pos = ltp.pos(hidden)
            getWords(len(sents))
        if '男' not in labels['noun'] and '女' not in labels['noun']:
            labels['noun'].append('男')
            labels['noun'].append('女')
        printAll()


def write2file():
    now = datetime.datetime.now()
    timing = now.strftime("%d-%H-%M-%S")
    res_filename = 'dataset/results/标注-' + timing + '.json'
    with open(res_filename, 'w', encoding='utf-8') as res:
        json_infos = json.dumps(all_infos, indent=2, ensure_ascii=False)
        res.write(json_infos)


def trans2text(filename):
    if not os.path.exists('dataset/text/'):
        os.mkdir('text', 0o755)
    if filename.endswith('.doc'):
        doc2text(filename)
        return True
    elif filename.endswith('.docx'):
        trans2txt(filename, 'text')
        return True
    elif filename.endswith('.txt'):
        fname = filename.split('/')[-1]
        new_filename = 'dataset/text/' + fname[0:fname.rfind('.')] + '.txt'
        try:
            shutil.copy(filename, new_filename)
        except shutil.SameFileError:
            pass
        return True
    return False


def getContent(filename):
    trans2text(filename)
    content = ''
    filename = filename.split('/')[-1]
    filename = 'dataset/text/' + filename[0:filename.rfind('.')] + '.txt'
    with open(filename, 'r', encoding='utf-8') as f:
        for c in f.read():
            content += c

    return content


def extractFromContent(content):
    getStopWords()
    getAccusations()
    global labels
    labels = {
        'noun': [],
        'verb': [],
        'adj': []
    }
    filename = 'case' + str(datetime.datetime.now()) + '.txt'
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(content)
    extract(filename)
    os.remove(filename)
    cleanLabels()
    # labels['noun'] = labels['noun'][:20]
    # labels['verb'] = labels['verb'][:20]
    # labels['adj'] = labels['adj'][:20]
    return cur_info, labels


def cleanLabels():
    global labels
    for i in range(len(labels['noun'])):
        labels['noun'][i] = labels['noun'][i].strip()
    for i in range(len(labels['verb'])):
        labels['verb'][i] = labels['verb'][i].strip()
    for i in range(len(labels['adj'])):
        labels['adj'][i] = labels['adj'][i].strip()
    labels['noun'] = list(set(labels['noun']))
    labels['verb'] = list(set(labels['verb']))[0:20]
    labels['adj'] = list(set(labels['adj']))


def saveAll(content, auto, labels):
    auto = json.loads(auto)
    labels = json.loads(labels)
    case_name = auto[0]['File']
    with open(case_name + '.txt', 'w', encoding='utf-8') as f:
        f.write(content)
    with open(case_name + '-手动标注.json', 'w', encoding='utf-8') as f:
        json_infos = json.dumps(labels, indent=2, ensure_ascii=False)
        f.write(json_infos)
    with open(case_name + '-自动标注.json', 'w', encoding='utf-8') as f:
        json_infos = json.dumps(auto, indent=2, ensure_ascii=False)
        f.write(json_infos)


if __name__ == '__main__':
    getStopWords()
    getAccusations()
    massExtract('dataset/text')
    write2file()
